package com.ruoyi;

import static org.opencv.highgui.HighGui.imshow;
import static org.opencv.highgui.HighGui.waitKey;
import static org.opencv.imgcodecs.Imgcodecs.imread;
import static org.opencv.imgcodecs.Imgcodecs.imwrite;
import org.opencv.core.*;
import org.opencv.imgproc.Imgproc;

import java.net.URL;
import java.util.*;

public class Test {
    public static void main(String[] args) {
        // 加载 OpenCV 库
        URL url = ClassLoader.getSystemResource("lib/opencv/opencv_java4110.dll");
        System.load(url.getPath());

        // 读取原始图像
        Mat rawImage = imread("C:/Data/System/Desktop/0.jpg");

        // Step 1: 全局白色平衡
        Mat balancedImage = whiteBalance(rawImage);
        imwrite("C:/Data/System/Desktop/1.jpg", balancedImage);

        // Step 2: 增强亮度大于指定阈值的“白”，其余保持不变
        Mat enhancedImage = enhanceBrightWhite(balancedImage, 200);
        imwrite("C:/Data/System/Desktop/2.jpg", enhancedImage);

        // Step 3: Lab 空间提取白色区域
        Mat whiteMask = extractWhiteRegionUsingLab(balancedImage);
        imwrite("C:/Data/System/Desktop/3.jpg", whiteMask);

        // Step 4: 使用自定义 erode + dilate 消融小点、保留大区域
        Mat kernel = Imgproc.getStructuringElement(Imgproc.MORPH_RECT, adaptiveKernelSize(whiteMask, 0.01));

        // Step 5: 腐蚀：去掉小点
        Imgproc.erode(whiteMask, whiteMask, kernel, new Point(-1, -1), 2);
        imwrite("C:/Data/System/Desktop/4.jpg", whiteMask);

        // Step 6: 膨胀：恢复主体区域
        Imgproc.dilate(whiteMask, whiteMask, kernel, new Point(-1, -1), 2);
        imwrite("C:/Data/System/Desktop/5.jpg", whiteMask);

        // Step 7: 可选：开运算，进一步去噪
        Imgproc.morphologyEx(whiteMask, whiteMask, Imgproc.MORPH_OPEN, kernel);
        imwrite("C:/Data/System/Desktop/6.jpg", whiteMask);

        // Step 8:
        List<MatOfPoint> contours = new ArrayList<>();
        Mat hierarchy = new Mat();
        Imgproc.findContours(whiteMask, contours, hierarchy, Imgproc.RETR_EXTERNAL, Imgproc.CHAIN_APPROX_SIMPLE);
        // 至少占图像面积的 30%
        double minArea = rawImage.rows() * rawImage.cols() * 0.3;
        MatOfPoint bestContour = null;
        double maxArea = minArea;
        for (MatOfPoint contour : contours) {
            double area = Imgproc.contourArea(contour);
            if (area < minArea) continue;

            MatOfPoint2f approxCurve = approximateContour(contour);
            if (approxCurve.total() == 4 && isValidQuadrilateral(approxCurve)) {
                if (area > maxArea) {
                    maxArea = area;
                    bestContour = new MatOfPoint(approxCurve.toArray());
                }
            }
        }

        // Step 7:
        if (bestContour != null) {
            // 绘制轮廓
            Mat result = rawImage.clone();
            Imgproc.drawContours(result, Collections.singletonList(bestContour), -1, new Scalar(0, 255, 0), 3);
            imwrite("C:/Data/System/Desktop/7.jpg", result);
            // 透视矫正
            Mat corrected = correctPerspective(rawImage, bestContour);
            imwrite("C:/Data/System/Desktop/8.jpg", corrected);
        } else {
            System.out.println("未找到合适的四边形");
        }
    }




    private static MatOfPoint2f approximateContour(MatOfPoint contour) {
        MatOfPoint2f curve = new MatOfPoint2f(contour.toArray());
        double epsilon = 0.02 * Imgproc.arcLength(curve, true);
        MatOfPoint2f approxCurve = new MatOfPoint2f();
        Imgproc.approxPolyDP(curve, approxCurve, epsilon, true);

        if (approxCurve.total() > 6) {
            epsilon += 0.01 * Imgproc.arcLength(curve, true);
            Imgproc.approxPolyDP(curve, approxCurve, epsilon, true);
        }

        return approxCurve;
    }

    private static boolean isValidQuadrilateral(MatOfPoint2f contour) {
        RotatedRect rect = Imgproc.minAreaRect(new MatOfPoint2f(contour.toArray()));
        double ratio = Math.min(rect.size.width, rect.size.height) /
                Math.max(rect.size.width, rect.size.height);
        return ratio > 0.1;
    }

    private static Mat correctPerspective(Mat src, MatOfPoint quad) {
        Point[] pts = quad.toArray();
        Point[] sortedPts = sortPointsClockwise(pts);
        // 计算四边形的两条边向量
        double edge1 = Math.sqrt(Math.pow(sortedPts[1].x - sortedPts[0].x, 2) + Math.pow(sortedPts[1].y - sortedPts[0].y, 2));
        double edge2 = Math.sqrt(Math.pow(sortedPts[3].x - sortedPts[0].x, 2) + Math.pow(sortedPts[3].y - sortedPts[0].y, 2));
        // 判断哪条边更长，作为宽度
        boolean isLandscape = edge1 > edge2;
        int n = 5;
        int width = isLandscape ? 297 * n : 210 * n;  // 如果是横版，交换宽高
        int height = isLandscape ? 210 * n : 297 * n;
        Mat srcPoints = new Mat(4, 1, CvType.CV_32FC2);
        for (int i = 0; i < 4; i++) {
            srcPoints.put(i, 0, new float[]{(float) sortedPts[i].x, (float) sortedPts[i].y});
        }
        Mat dstPoints = new Mat(4, 1, CvType.CV_32FC2);
        dstPoints.put(0, 0, new float[]{0, 0});
        dstPoints.put(1, 0, new float[]{width, 0});
        dstPoints.put(2, 0, new float[]{width, height});
        dstPoints.put(3, 0, new float[]{0, height});
        Mat transform = Imgproc.getPerspectiveTransform(srcPoints, dstPoints);
        Mat corrected = new Mat();
        Imgproc.warpPerspective(src, corrected, transform, new Size(width, height));
        return corrected;
    }

    private static Point[] sortPointsClockwise(Point[] pts) {
        Point center = new Point(0, 0);
        for (Point p : pts) {
            center.x += p.x / pts.length;
            center.y += p.y / pts.length;
        }
        List<Point> list = new ArrayList<>();
        Collections.addAll(list, pts);
        list.sort(Comparator.comparingDouble(p -> Math.atan2(p.y - center.y, p.x - center.x)));
        return list.toArray(new Point[0]);
    }

    // Step 1: 全局白平衡
    private static Mat whiteBalance(Mat src) {
        Mat dst = new Mat();
        Core.normalize(src, dst, 0, 255, Core.NORM_MINMAX);
        List<Mat> channels = new ArrayList<>();
        Core.split(dst, channels);
        double avgR = Core.mean(channels.get(2)).val[0];
        double avgG = Core.mean(channels.get(1)).val[0];
        double avgB = Core.mean(channels.get(0)).val[0];

        double avgGray = (avgR + avgG + avgB) / 3.0;

        for (int i = 0; i < channels.size(); i++) {
            Scalar scale = new Scalar(avgGray / (Core.mean(channels.get(i)).val[0]));
            Core.multiply(channels.get(i), scale, channels.get(i));
            Core.MinMaxLocResult result = Core.minMaxLoc(channels.get(i));
            double min = result.minVal;
            double max = result.maxVal;
            if (max != min) {
                Core.normalize(channels.get(i), channels.get(i), 0, 255, Core.NORM_MINMAX);
            }
        }

        Core.merge(channels, dst);
        return dst;
    }

    // Step 2: 增强亮度大于指定阈值的“白”，其余保持不变
    private static Mat enhanceBrightWhite(Mat src, int threshold) {
        // Step 1: 转换为灰度图用于判断哪些是“高亮度”的白
        Mat gray = new Mat();
        Imgproc.cvtColor(src, gray, Imgproc.COLOR_BGR2GRAY);

        // Step 2: 创建掩膜：亮度大于 threshold 的像素设为 255（即选中）
        Mat mask = new Mat();
        Core.inRange(gray, new Scalar(threshold), new Scalar(255), mask);

        // Step 3: 创建增强后的图像（每个通道 + 偏移量）
        Mat enhanced = new Mat();
        Core.add(src, Scalar.all(55), enhanced);  // 将所有像素加 55

        // Step 4: 防止溢出
        Imgproc.threshold(enhanced, enhanced, 255, 255, Imgproc.THRESH_TRUNC);

        // Step 5: 创建结果图像，初始化为原始图像
        Mat result = src.clone();

        // Step 6: 使用双重循环，只增强 mask 中为 255 的像素
        for (int row = 0; row < mask.rows(); row++) {
            for (int col = 0; col < mask.cols(); col++) {
                byte[] maskData = new byte[1];
                mask.get(row, col, maskData);
                if (maskData[0] == (byte) 255) {
                    // 获取增强后的像素值（假设是三通道 CV_8UC3 图像）
                    byte[] pixelData = new byte[3];
                    enhanced.get(row, col, pixelData);

                    // 转换为 double 数组写回 result（也可以直接操作 byte）
                    double[] outPixel = new double[] {
                            pixelData[0] & 0xFF,
                            pixelData[1] & 0xFF,
                            pixelData[2] & 0xFF
                    };
                    result.put(row, col, outPixel);
                }
            }
        }

        return result;
    }

    // Step 3: Lab 空间提取白色区域
    private static Mat extractWhiteRegionUsingLab(Mat src) {
        Mat lab = new Mat();
        Imgproc.cvtColor(src, lab, Imgproc.COLOR_BGR2Lab);

        List<Mat> channels = new ArrayList<>();
        Core.split(lab, channels);
        Mat L = channels.get(0);  // 亮度
        Mat A = channels.get(1);  // 绿红轴
        Mat B = channels.get(2);  // 蓝黄轴

        // 定义 a 和 b 的阈值范围（越小越接近白色）
        int thresholdAB = 30;

        Mat maskAB = new Mat();
        Mat absA = new Mat(), absB = new Mat();
        Core.absdiff(A, Scalar.all(128), absA);  // Lab 中 a=128 表示无色差
        Core.absdiff(B, Scalar.all(128), absB);  // 同理 b=128

        // 创建 a 和 b 的掩膜
        Core.inRange(absA, Scalar.all(0), Scalar.all(thresholdAB), maskAB);
        Mat maskB = new Mat();
        Core.inRange(absB, Scalar.all(0), Scalar.all(thresholdAB), maskB);
        Core.bitwise_and(maskAB, maskB, maskAB);

        // 可选：结合亮度约束（不要太暗也不要太亮）
        Mat maskL = new Mat();
        Imgproc.threshold(L, maskL, 200, 255, Imgproc.THRESH_BINARY); // 高亮度部分

        Mat finalMask = new Mat();
        Core.bitwise_and(maskAB, maskL, finalMask);

        return finalMask;
    }

    // 根据图像大小自动计算合适的结构元素大小
    public static Size adaptiveKernelSize(Mat image, double ratio) {
        int size = (int) Math.max(3, Math.round(image.cols() * ratio));
        return new Size(size, size);
    }
}